from matplotlib.pyplot import figure, draw, show
from matplotlib.animation import FuncAnimation
from mayavi import mlab
from wx import Yield
from cylinder import sg_plot, set_sg_data
from pdb import set_trace
"""
Module to create assorted plots and graphics of the AFD runs
"""

def temperature(self,t,T):
    """
    Create a plot of the temperatures
    t - [lx1] numpy array time
    T - [lxnxm] numpy array temperature
    return matplotlib.figure
    """
    T = T - 273
    t = t/3600.
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(t,np.average(np.average(T,axis=2),axis=1),label=r'$T_{\mathrm{average}}$',color='blue',ls='-')
    ax.plot(t,T[:,T.shape[1]/2,T.shape[2]/3],label=r'$T_{r=0,z=.3H}$',color='green',ls='--')
    ax.plot(t,T[:,T.shape[1]/2,T.shape[2]/2],label=r'$T_{r=0,z=.5H}$',color='red',ls=':')
    ax.plot(t,T[:,T.shape[1]/2,2*T.shape[2]/3],label=r'$T_{r=0,z=.7H}$',color='cyan',ls='-.')
    ax.set_xlabel(r'$t $(hr)')
    ax.set_ylabel(r'$T (^{o}\mathrm{C})$')
    ax.legend(loc=4)
    return fig

def solvent_conc(self,t,w_L):
    """
    Create a plot of the solvent concentration
    t - [lx1] numpy array time (sec)
    w_L - [lxnxm] numpy array solvent weight fraction
    """
    t = t/3600.
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(t,np.average(np.average(w_L,axis=2),axis=1),label=r'$w_{\mathrm{average}}$',color='blue',ls='-')
    ax.plot(t,w_L[:,w_L.shape[1]/2,w_L.shape[2]/3],label=r'$w_{r=0,z=.3H}$',color='green',ls='--')
    ax.plot(t,w_L[:,w_L.shape[1]/2,w_L.shape[2]/2],label=r'$w_{r=0,z=.5H}$',color='red',ls=':')
    ax.plot(t,w_L[:,w_L.shape[1]/2,2*w_L.shape[2]/3],label=r'$w_{r=0,z=.7H}$',color='cyan',ls='-.')
    ax.set_xlabel(r'$t $(hr)')
    ax.set_ylabel(r'$w_L (^w / _w)$')
    ax.legend(loc=1)
    return fig

def animated_plots(self,r,z,v_L,T,fname,anim=False):
    """
    Create animated contour plots of the temperature and solvent concentration
    r - [nxm] numpy array radial nodal points 
    z - [nxm] numpy array axial nodal points
    v_L - [lxnxm] numpy array solvent concentration
    T - [lxnxm] numpy array temperature
    fname - string filename of movie
    anim - boolean use matplotlib's animation module
    """
    #Nt = v_L.shape[0]
    #NR = v_L.shape[1]
    #NH = v_L.shape[2]
    #z,r = np.meshgrid(np.linspace(1e-2*self.dz,H-1e-2*self.dz,NH),np.linspace(1e-2*self.dr,R-1e-2*self.dr,NR))
    if anim == True:
        fig = plt.figure()
        ims = []
        from matplotlib import animation
        for i in xrange(0,T.shape[0],T.shape[0]/200):
            im = plt.contourf(r,z,T[i])
            def setvisible(self,vis):
               for c in self.collections: c.set_visible(vis)
            im.set_visible = types.MethodType(setvisible,im,None)
            im.axes = plt.gca()
            im.figure = fig 
            ims.append([im])
        ani = animation.ArtistAnimation(fig, ims, interval=50, blit=False, repeat_delay=1000)
        ani.save('dynamic_contour.mp4')
        plt.show()
    else:
        cnt = 0
        fig = plt.figure()
        ax1 = fig.add_subplot(121)
        ax2 = fig.add_subplot(122)
        for i in xrange(0,T.shape[0],T.shape[0]/200):
            ax1.contourf(r,z,T[i],np.linspace(np.min(T)-1,np.max(T)+1))
            ax2.contourf(r,z,v_L[i],np.linspace(np.min(v_L)-.01,np.max(v_L)+0.01))
            ax1.set_xlim(r[0,0],r[-1,0])
            ax1.set_ylim(z[0,0],z[0,-1])
            ax2.set_xlim(r[0,0],r[-1,0])
            ax2.set_ylim(z[0,0],z[0,-1])
            ax1.set_title('Temperature (K)')
            ax2.set_title(r'Solvent Concentration ($\phi_L$)')
            plt.savefig('%s_%03d.png'%(fname,cnt))
            plt.cla()
            cnt += 1
        command = ('ffmpeg','-r','10','-b','1800','-i','%s_'%(fname)+'%03d.png','%s.mp4'%fname)
        subprocess.check_call(command)
    return

def composite_figure(r, z, T, vL):
    """
    Create a composite figure of the simulation data.
    r : [nr x nz] numpy array of the radial coordinates
    z : [nr x nz] numpy array of the axial coordinates
    T : [nt x nr x nz] numpy array of temperature data
    vL : [nt x nr x nz] numpy array of liquid volume fraction data
    """
    def cylinder(data, name):
        f = sg_plot(r, z, 25, data[0], name)
        d = f.children[0]
        sg = d.data
        for i, dat in enumerate(data[0:-1:1000]):
            sg = set_sg_data(sg, dat, name)
            d.update()
            yield mlab.screenshot(figure=f)

    fig = figure()
    draw()
    show()
    
    ax_T = fig.add_subplot(1,2,1)
    ax_vL = fig.add_subplot(1,2,2)
    anim = FuncAnimation(fig, cylinder, fargs=None, save_count=None, **kwargs)
    for Ti in gen_T:
        vLi = gen_vL.next()
        ax_T.imshow(Ti)
        ax_vL.imshow(vLi)
        draw()
        show()

if __name__=="__main__":
    pass
